package com.site.common.aspect;

import com.google.common.collect.Maps;
import com.google.common.util.concurrent.RateLimiter;
import com.site.common.annotation.AccessLimit;
import com.site.common.annotation.enums.AccessLimitResourceType;
import com.site.common.annotation.enums.AccessLimitStoreType;
import com.site.common.annotation.enums.AccessLimitType;
import com.site.common.exception.BusinessException;
import com.site.util.RedisUtil;
import org.apache.commons.lang3.StringUtils;
import org.aspectj.lang.ProceedingJoinPoint;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.util.ObjectUtils;

import javax.servlet.http.HttpServletRequest;
import java.lang.reflect.Method;
import java.net.InetAddress;
import java.net.UnknownHostException;
import java.text.SimpleDateFormat;
import java.util.Arrays;
import java.util.Collections;
import java.util.Date;
import java.util.Enumeration;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;

/** 弃用了，使用限流过滤器
 * 限流切面
 */
public class RateLimiterAspect extends AbstractAspect {

    private static Logger logger = LoggerFactory.getLogger(RateLimiterAspect.class);

    //添加速率.保证是单例的
    private static RateLimiter rateLimiter = RateLimiter.create(1000);
    //使用url做为key,存放令牌桶 防止每次重新创建令牌桶
    private static Map<String, RateLimiter> limiterMap = Maps.newConcurrentMap();

    @Override
    protected Object execute(ProceedingJoinPoint joinPoint, Method method, HttpServletRequest request) {
        AccessLimit accessLimit = method.getAnnotation(AccessLimit.class);
        String key = getKey(request,accessLimit,method);
        AccessLimitStoreType storeType = accessLimit.store();
        // guava
        if (Objects.equals(storeType,AccessLimitStoreType.GUAVA)){
            // 接口限流
            if (Objects.equals(accessLimit.type(),AccessLimitType.LIMIT)){
                return limitGuava(key,accessLimit);
            }
            // 限制访问次数
            //if (Objects.equals(accessLimit.type(),AccessLimitType.COUNT)){
            //}
        }
        // redis
        if (Objects.equals(storeType,AccessLimitStoreType.REDIS)){
            int limitCount = accessLimit.count();
            key = accessLimit.prefix() + key;
            try {
                // 限流
                if (Objects.equals(accessLimit.type(),AccessLimitType.LIMIT)){
                    return limitRedis(joinPoint,Collections.singletonList(key),limitCount,accessLimit.timeOut()==0?600:accessLimit.timeOut(),accessLimit.timeOutUnit());
                }
                // 限制访问次数
                if (Objects.equals(accessLimit.type(),AccessLimitType.COUNT)){
                    //ImmutableList<String> keys = ImmutableList.of(StringUtils.join(accessLimit.prefix(),key));
                    int timeOut = accessLimit.timeOut();
                    return limitRedisCount(joinPoint,key,Collections.singletonList(key),limitCount,Objects.equals(timeOut,0)?(3600*24*1000):timeOut,accessLimit.timeOutUnit());
                }
            }catch (Throwable e){
                if(e instanceof BusinessException){
                    throw new BusinessException(e.getLocalizedMessage());
                }
                throw new RuntimeException(e);
            }
        }
        return null;
    }

    /**
     * guava限流
     * @param key
     * @param accessLimit
     * @return
     */
    private Object limitGuava(String key,AccessLimit accessLimit){
        if (!limiterMap.containsKey(key)){
            // 创建令牌桶
            rateLimiter = RateLimiter.create(accessLimit.number());
            limiterMap.put(key,rateLimiter);
            logger.info("<<=================  请求地址{},创建令牌桶,容量{} 成功!!!",key,accessLimit.number());
        }
        rateLimiter = limiterMap.get(key);
        if (!rateLimiter.tryAcquire(accessLimit.timeOut(),accessLimit.timeOutUnit())){//获取令牌
            logger.info("Error ---时间:{},获取令牌失败.", new SimpleDateFormat("yyyy-MM-dd HH:mm:ss").format(new Date()));
            throw new BusinessException("服务器繁忙，请稍后再试!");
        }
        return null;
    }

    /**
     * redis 限流
     * @param keys
     * @param limitCount
     * @param timeOut
     * @param timeUnit
     * @return
     */
    private Object limitRedis(ProceedingJoinPoint joinPoint,List<String> keys, int limitCount, int timeOut, TimeUnit timeUnit) throws Throwable {
        /*Object count = RedisUtil.instance.getRedisTemplate().opsForValue().get(key);
        if (null == count) {
            RedisUtil.instance.getRedisTemplate().opsForValue().set(key,1,timeOut,timeUnit);
            return null;
        }
        if (Integer.parseInt(String.valueOf(count)) >= limitCount){
            throw new BusinessException("服务器繁忙，请稍后再试!");
        }
        RedisUtil.instance.getRedisTemplate().opsForValue().increment(key);
        return null;*/
        String luaScript = buildLuaScript();
        RedisScript<Number> redisScript = new DefaultRedisScript<>(luaScript, Number.class);
        // redisScript.setScriptSource(new ResourceScriptSource(new ClassPathResource("concurrent_request_rate_limiter.lua")));
        // 如果执行lua脚本报错,请转字符串调用,Number count = RedisUtil.instance.getRedisTemplate().execute(redisScript, keys, String.valueOf(limitCount), String.valueOf(timeOut / 1000),timeUnit);// timeOut/1000 是因为默认为秒
        Number count = RedisUtil.instance.getRedisTemplate().execute(redisScript, keys, limitCount, timeOut / 1000,timeUnit);// timeOut/1000 是因为默认为秒
        //logger.info("Access try count is {} and key = {}", count, key);
        if(count !=null && count.intValue() <= limitCount){
            return joinPoint.proceed();
        }else{
            throw new BusinessException("服务器繁忙，请稍后再试!");
        }
    }

    /**
     * redis 限制访问次数
     * @param joinPoint
     * @param key
     * @param keys
     * @param limitCount
     * @param timeOut
     * @return
     * @throws Throwable
     */
    private Object limitRedisCount(ProceedingJoinPoint joinPoint, String key, List<String> keys, int limitCount, int timeOut, TimeUnit timeUnit) throws Throwable {
        String luaScript = buildLuaScript();
        RedisScript<Number> redisScript = new DefaultRedisScript<>(luaScript, Number.class);
        Number count = RedisUtil.instance.getRedisTemplate().execute(redisScript, keys, limitCount, timeOut / 1000,timeUnit);
        logger.info("Access try count is {} and key = {}", count, key);
        if(count !=null && count.intValue() <= limitCount){
            return joinPoint.proceed();
        }else{
            throw new BusinessException("访问超限");
        }
    }

    /**
     * 获取key
     * @param request
     * @param accessLimit
     * @param method
     * @return
     */
    private String getKey(HttpServletRequest request,AccessLimit accessLimit,Method method){
        AccessLimitResourceType limitResourceType = accessLimit.resourceType();
        //如果类型是IP，则根据IP限制访问次数，key取IP地址
        if (AccessLimitResourceType.IP.equals(limitResourceType)){
            return getIpAddr(request) + request.getRequestURI();
        }
        //如果类型是customer，则根据key限制访问次数
        if (AccessLimitResourceType.CUSTOMER.equals(limitResourceType)){
            //return accessLimit.key();
            return getSpelValueByKey(accessLimit.key(), request);
        }
        //否则按照方法名称限制访问次数
        return method.getName();
    }

    /**
     * lua限流脚本
     * @return
     */
    public String buildLuaScript(){
        StringBuilder sb = new StringBuilder();
        //定义c
        sb.append("local c");
        //获取redis中的值
        sb.append("\nc = redis.call('get',KEYS[1])");
        //如果调用不超过最大值
        sb.append("\nif c and tonumber(c) > tonumber(ARGV[1]) then");
        //直接返回
        sb.append("\n return c;");
        //结束
        sb.append("\nend");
        //访问次数加一
        sb.append("\nc = redis.call('incr',KEYS[1])");
        //如果是第一次调用
        sb.append("\nif tonumber(c) == 1 then");
        //设置对应值的过期设置
        sb.append("\nredis.call('expire',KEYS[1],ARGV[2])");
        //结束
        sb.append("\nend");
        //返回
        sb.append("\nreturn c;");
        return sb.toString();
    }

    private static final String IP_UTILS_FLAG = ",";
    private static final String UNKNOWN = "unknown";
    private static final String LOCALHOST_IP = "0:0:0:0:0:0:0:1";
    private static final String LOCALHOST_IP1 = "127.0.0.1";
    /**
     * 获取IP地址
     * <p>
     * 使用Nginx等反向代理软件， 则不能通过request.getRemoteAddr()获取IP地址
     * 如果使用了多级反向代理的话，X-Forwarded-For的值并不止一个，而是一串IP地址，X-Forwarded-For中第一个非unknown的有效IP字符串，则为真实IP地址
     */
    public static String getIpAddr(HttpServletRequest request) {
        String ip = null;
        try {
            //以下两个获取在k8s中，将真实的客户端IP，放到了x-Original-Forwarded-For。而将WAF的回源地址放到了 x-Forwarded-For了。
            ip = request.getHeader("X-Original-Forwarded-For");
            if (StringUtils.isEmpty(ip) || UNKNOWN.equalsIgnoreCase(ip)) {
                ip = request.getHeader("X-Forwarded-For");
            }
            //获取nginx等代理的ip
            if (StringUtils.isEmpty(ip) || UNKNOWN.equalsIgnoreCase(ip)) {
                ip = request.getHeader("x-forwarded-for");
            }
            if (StringUtils.isEmpty(ip) || UNKNOWN.equalsIgnoreCase(ip)) {
                ip = request.getHeader("Proxy-Client-IP");
            }
            if (StringUtils.isEmpty(ip) || ip.length() == 0 || UNKNOWN.equalsIgnoreCase(ip)) {
                ip = request.getHeader("WL-Proxy-Client-IP");
            }
            if (StringUtils.isEmpty(ip) || UNKNOWN.equalsIgnoreCase(ip)) {
                ip = request.getHeader("HTTP_CLIENT_IP");
            }
            if (StringUtils.isEmpty(ip) || UNKNOWN.equalsIgnoreCase(ip)) {
                ip = request.getHeader("HTTP_X_FORWARDED_FOR");
            }
            //兼容k8s集群获取ip
            if (StringUtils.isEmpty(ip) || UNKNOWN.equalsIgnoreCase(ip)) {
                ip = request.getRemoteAddr();
                if (LOCALHOST_IP1.equalsIgnoreCase(ip) || LOCALHOST_IP.equalsIgnoreCase(ip)) {
                    //根据网卡取本机配置的IP
                    InetAddress iNet = null;
                    try {
                        iNet = InetAddress.getLocalHost();
                    } catch (UnknownHostException e) {
                        logger.error("getClientIp error: {}", e);
                    }
                    ip = iNet.getHostAddress();
                }
            }
        } catch (Exception e) {
            logger.error("IPUtils ERROR ", e);
        }
        //使用代理，则获取第一个IP地址
        if (!StringUtils.isEmpty(ip) && ip.indexOf(IP_UTILS_FLAG) > 0) {
            ip = ip.substring(0, ip.indexOf(IP_UTILS_FLAG));
        }
        return ip;
    }

    /**
     * 获取参数(注意：获取不到json参数，只能获取表单参数,如:?name=小二&age=18)
     * @param request
     * @return
     */
    public String getParameter(HttpServletRequest request){
        StringBuilder sb = new StringBuilder(100);
        Enumeration<String> parameteres = request.getParameterNames();
        if (parameteres.hasMoreElements()){
            while (parameteres.hasMoreElements()){
                String key = parameteres.nextElement();
                String[] values = request.getParameterValues(key);
                sb.append(key).append("=");
                for (String val : values){
                    sb.append(val);
                }
                sb.append(",");
            }
        }
        return sb.toString().length() > 0 ? sb.toString().substring(0,sb.length()-1) : "";
    }

    /**
     * 获取spring EL表达式的key
     * @return
     */
    public String getSpelKey(String spelKey){
        if (null == spelKey || spelKey.length() == 0 || !spelKey.contains("#")) {
            return null;
        }
        return spelKey.replaceAll("#\\{","").replaceAll("}","");
    }

    /**
     * 获取spring EL表达式的key的值
     * @return
     */
    public String getSpelValueByKey(String spelKey,HttpServletRequest request){
        String parameter = this.getParameter(request);
        if (!ObjectUtils.isEmpty(parameter)){
            Map<String, String> parameterMap = Arrays.stream(parameter.split(",")).map(s -> s.split("=")).collect(Collectors.toMap(item -> item[0], item -> item[1]));
            if (!ObjectUtils.isEmpty(parameterMap)){
                return parameterMap.get(getSpelKey(spelKey));
            }
        }
        return null;
    }
}
